package middleware

import (
	"context"
	"net/http"
	"time"
)

type TimeOut struct {
	Next http.Handler
}

func (t *TimeOut) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
	if t.Next == nil {
		t.Next = http.DefaultServeMux
	}

	// 替换context
	ctx := request.Context()
	// 生成新context
	ctx, cancelFunc := context.WithTimeout(ctx, time.Second*3)
	defer func() { cancelFunc() }()

	// 替换
	request.WithContext(ctx)
	// 正常结束信号通道
	signalChan := make(chan struct{}, 1)
	go func() {
		t.Next.ServeHTTP(writer, request)
		signalChan <- struct{}{}
	}()
	select {
	case <-signalChan:
		// 收到正常结束信号
		return
	case <-ctx.Done():
		// 超时
		http.Error(writer, http.StatusText(http.StatusRequestTimeout),
			http.StatusRequestTimeout)
	}
	ctx.Done()
}
